#include <inttypes.h>
#include <math.h>
#include <pthread.h>
#include <signal.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include "common/getopt.h"
#include "common/timestamp.h"
#include "lcm/lcm.h"
#include "lcmtypes/battery_t.h"
#include "lcmtypes/diff_drive_t.h"
#include "lcmtypes/motor_feedback_t.h"
#include "uorc/uorc.h"

//////////////////////////////////////////////////////////////////////
//////////////////////////////// Structs /////////////////////////////
//////////////////////////////////////////////////////////////////////
typedef struct message_wrapper message_wrapper_t;
struct message_wrapper
{
    void *msg;
    pthread_mutex_t mutex;
};

typedef struct motordriver_data motordriver_data_t;
struct motordriver_data
{
    int num_threads;
    pthread_t *pids;

    lcm_t *lcm;
    char *channel_prefix;
    int hz_feedback;
    int hz_battery;

    message_wrapper_t msgmanual;
    message_wrapper_t msgauto;
    message_wrapper_t msgencoder;

    uorc_t *uorc;
    uorc_motor_t *motor_right;
    uorc_motor_t *motor_left;
    uorc_encoder_t *encoder_right;
    uorc_encoder_t *encoder_left;

    int quit;
    int verbose;
};

//////////////////////////////////////////////////////////////////////
////////////////////////// Helper Functions //////////////////////////
//////////////////////////////////////////////////////////////////////
static void cleanup(int s, void *arg)
{
    motordriver_data_t *md = (motordriver_data_t *)arg;
    md->quit = 1;

    // join
    for (int i = 0; i < md->num_threads; i++) {
        pthread_join(md->pids[i], NULL);
    }
    if (md->uorc != NULL){
        uorc_destroy(md->uorc);
    }
    if (md->channel_prefix != NULL) {
        free(md->channel_prefix);
    }

    pthread_mutex_destroy(&md->msgauto.mutex);
    pthread_mutex_destroy(&md->msgencoder.mutex);
    pthread_mutex_destroy(&md->msgmanual.mutex);

    if (md->msgauto.msg != NULL) {
        free(md->msgauto.msg);
    }
    if (md->msgencoder.msg != NULL) {
        free(md->msgencoder.msg);
    }
    if (md->msgmanual.msg != NULL) {
        free(md->msgmanual.msg);
    }

    if (md->lcm != NULL) {
        lcm_destroy(md->lcm);
    }

    if (md->pids != NULL)
        free(md->pids);
    md->pids = NULL;

    printf("done\n");
}

static void sighandler(int s)
{
	switch(s) {
        case SIGSEGV:
        case SIGINT:
        case SIGTERM:
            printf("Exiting...");
            exit(s);
        default:
            break;
	}
}

float max_float(float a, float b)
{
    return a > b ? a : b;
}

void compute_motor_command(float *cmd)
{
    // consider velocities less than this to be zero.
    const float DEADBAND_THRESH = 0.05;

    float minpwm = 0.2;  // ff pwm corresponding to speed = 0
    float maxpwm = 1.0;  // ff pwm corresponding to speed = 1

    for (int i = 0;i < 2; i++) {
        float goal = cmd[i];
        float abs_goal = fabs(goal);
        float ffpwm;
        if (abs_goal < DEADBAND_THRESH)
            ffpwm = 0;
        else {
            ffpwm = minpwm + (maxpwm - minpwm)*abs_goal;
            if (goal < 0)
                ffpwm = -ffpwm;
        }
        cmd[i] = ffpwm;
    }
}

/* return channel name as: <prefix>_<base> */
char * make_channelname(const char *prefix, const char *base)
{
    char *channel = NULL;
    int prefix_len = 0;

    if (prefix != NULL && base != NULL && (prefix_len = (int)strlen(prefix))) {
        channel = (char *) calloc(prefix_len + strlen(base) + 2, sizeof(char));
        strcpy(channel, prefix);
        channel[prefix_len] = '_';
        channel[prefix_len+1] = '\0';
        strcat(channel, base);
    } else
        channel = strdup(base);
    return channel;
}

//////////////////////////////////////////////////////////////////////
///////////////////////////// LCM HANDLERS ///////////////////////////
//////////////////////////////////////////////////////////////////////
static void handler_diffdrive(const lcm_recv_buf_t *rbuf, const char *channel,
                              const diff_drive_t *msg, void *user)
{
    message_wrapper_t *data = user;

    pthread_mutex_lock(&data->mutex);
    {
        if (data->msg)
            free(data->msg);
        data->msg = diff_drive_t_copy(msg);
    }
    pthread_mutex_unlock(&data->mutex);
}


//////////////////////////////////////////////////////////////////////
/////////////////////////////// THREADS //////////////////////////////
//////////////////////////////////////////////////////////////////////

// motor thread
void *run_motor_thread(void *v)
{
    const int64_t TIMEOUT_MANUAL_US = 250000;
    const int64_t TIMEOUT_DIFF_US   = 250000;

    motordriver_data_t *md = (motordriver_data_t *)v;

    diff_drive_t actual_motor_command;
    actual_motor_command.left_enabled = 1;
    actual_motor_command.right_enabled = 1;

    float cmdLR[2];

    ///// LCM Subscribe /////
    char *channel = make_channelname(md->channel_prefix, "DIFF_DRIVE");
    printf("NFO: subscribing to: %s\n", channel);
    diff_drive_t_subscription_t *subscription =
        diff_drive_t_subscribe(md->lcm, channel, handler_diffdrive, (void *)&md->msgauto);

    char *channel_publish = make_channelname(md->channel_prefix, "MOTOR_COMMAND");
    printf("NFO: publishing motor commands on: %s\n", channel_publish);

    int print_counter = 0;
	while (!md->quit) {

        // stop if no valid command received
        cmdLR[0] = 0.0;
        cmdLR[1] = 0.0;

        int64_t utime = timestamp_now();

        char *cmdName = "NONE";

        // autonomous mode
        pthread_mutex_lock(&md->msgauto.mutex);
        {
            diff_drive_t *autonomous = (diff_drive_t *)md->msgauto.msg;
            if (autonomous != NULL && (utime - autonomous->utime < TIMEOUT_DIFF_US)) {
                cmdName = "AUTO";
                cmdLR[0] = autonomous->left;
                cmdLR[1] = autonomous->right;
            }
        }
        pthread_mutex_unlock(&md->msgauto.mutex);

        // manual mode
        pthread_mutex_lock(&md->msgmanual.mutex);
        {
            diff_drive_t *manual = (diff_drive_t *)md->msgmanual.msg;
            if (manual != NULL && (utime - manual->utime < TIMEOUT_MANUAL_US)) {
                cmdName = "MANUAL";
                cmdLR[0] = manual->left;
                cmdLR[1] = manual->right;
            }
        }
        pthread_mutex_unlock(&md->msgmanual.mutex);

        if ((print_counter++ % 20) == 0)
            printf("%s\t%3.3f\t%3.3f\n", cmdName, cmdLR[0], cmdLR[1]);

        // ensure commands in [-1, 1]
        float max = max_float(fabs(cmdLR[0]), fabs(cmdLR[1]));
        if (max > 1.0) {
            cmdLR[0] /= max;
            cmdLR[1] /= max;
        }

        compute_motor_command(cmdLR);

        uorc_motor_set_pwm(md->motor_left, (float)cmdLR[0]);
        uorc_motor_set_pwm(md->motor_right, (float)cmdLR[1]);

        // publish command (nothing subscribes to this, unless logging)
        actual_motor_command.utime = utime;
        actual_motor_command.left = cmdLR[0];
        actual_motor_command.right = cmdLR[1];

        diff_drive_t_publish(md->lcm, channel_publish, &actual_motor_command);

        usleep(50000);
    }
    free(channel);
    free(channel_publish);
    diff_drive_t_unsubscribe(md->lcm, subscription);
    pthread_exit(NULL);
}

// encoder thread
void *run_encoder_thread(void *v)
{
    motordriver_data_t *md = (motordriver_data_t *)v;
    motor_feedback_t mf;
    mf.nmotors = 2;
    mf.encoders = (int32_t *) calloc(mf.nmotors, sizeof(int32_t));
    mf.current = (float *) calloc(mf.nmotors, sizeof(float));
    mf.applied_voltage = (float *) calloc(mf.nmotors, sizeof(float));
    uorc_status_t status;

    int period = (int) round(1000000.0 / md->hz_feedback);
    char *channel = make_channelname(md->channel_prefix, "MOTOR_FEEDBACK");

    printf("NFO: encoder publishing at %d Hz (period = %d)\n", md->hz_feedback, period);

	while (!md->quit) {
        uorc_get_status(md->uorc, &status);

        float currentLeft = uorc_motor_get_current(md->motor_left, &status);
        float currentRight = uorc_motor_get_current(md->motor_right, &status);

        int32_t encoderLeftpos  = uorc_encoder_get_position(md->encoder_left,  &status);
        int32_t encoderRightpos = uorc_encoder_get_position(md->encoder_right, &status);

        float voltage = uorc_get_battery_voltage(&status);
        float voltageLeft  = voltage * uorc_motor_get_pwm(md->motor_left,  &status);
        float voltageRight = voltage * uorc_motor_get_pwm(md->motor_right, &status);

        mf.utime = timestamp_now();
        mf.encoders[0] = encoderLeftpos;
        mf.encoders[1] = encoderRightpos;

        mf.current[0] = currentLeft;
        mf.current[1] = currentRight;

        mf.applied_voltage[0] = voltageLeft;
        mf.applied_voltage[1] = voltageRight;

        mf.estop = uorc_get_estop_state(&status);

        motor_feedback_t_publish(md->lcm, channel, &mf);

        usleep(period);
    }
    free(channel);
    free(mf.encoders);
    free(mf.current);
    free(mf.applied_voltage);
    pthread_exit(NULL);
}

// battery thread
void *run_battery_thread(void *v)
{
    motordriver_data_t *md = (motordriver_data_t *)v;

    battery_t bt;
    uorc_status_t status;

    int period = (int) round(1000000.0 / md->hz_battery);
    char *channel = make_channelname(md->channel_prefix, "BATTERY");

    printf("NFO: publishing battery voltage at %d Hz (period = %d)\n", md->hz_battery, period);

	while (!md->quit) {
        uorc_get_status(md->uorc, &status);

        bt.utime = timestamp_now();
        bt.voltage = uorc_get_battery_voltage(&status);

        battery_t_publish(md->lcm, channel, &bt);
        usleep(period);
    }
    free(channel);
    pthread_exit(NULL);
}

void start_threads(motordriver_data_t *md)
{
    if (pthread_create(&(md->pids[0]), NULL, (void *) run_motor_thread, md)) {
        printf("ERR: %s: Failed to start thread\n", "motor_thread");
        exit(-1);
    }

    if (pthread_create(&(md->pids[1]), NULL, (void *) run_encoder_thread, md)) {
        printf("ERR: %s: Failed to start thread\n", "encoder_thread");
        exit(-1);
    }

    if (pthread_create(&(md->pids[2]), NULL, (void *) run_battery_thread, md)) {
        printf("ERR: %s: Failed to start thread\n", "battery_thread");
        exit(-1);
    }
}

//////////////////////////////////////////////////////////////////////
////////////////////////////////// MAIN //////////////////////////////
//////////////////////////////////////////////////////////////////////
int main(int argc, char *argv[])
{
    motordriver_data_t md = { .num_threads = 3, .quit = 0, .verbose = 0,
                              .uorc = NULL, .pids = NULL,
                              .motor_right = NULL, .motor_left = NULL,
                              .encoder_right = NULL, .encoder_left = NULL };

    getopt_t *gopt = getopt_create();
    getopt_add_string(gopt, 'p', "prefix", "",
                      "LCM channel prefix (<prefix>_DIFF_DRIVE and <prefix>_MOTOR_FEEDBACK)");
    getopt_add_int(gopt, 'f', "hz_feedback", "20", "<prefix>_MOTOR_FEEDBACK publish frequency");
    getopt_add_int(gopt, 'b', "hz_battery", "1", "<prefix>_BATTERY publish frequency");
    getopt_add_spacer(gopt, "s");
    getopt_add_bool(gopt, 'v', "verbose", 0, "Verbose output");

    if (!getopt_parse(gopt, argc, argv, 1)) {
        printf("motordriver options:\n");
        getopt_do_usage(gopt);
        return 0;
    }
    md.verbose = getopt_get_bool(gopt, "verbose");
    md.hz_feedback = getopt_get_int(gopt, "hz_feedback");
    md.hz_battery  = getopt_get_int(gopt, "hz_battery");

    ///// LCM /////
    md.lcm = lcm_create(NULL); // local traffic only
    if (md.lcm == NULL) {
        printf("ERR: Could not create lcm\nExiting...\n");
        exit(-1);
    }
    md.channel_prefix = getopt_get_string(gopt, "prefix");

    ///// uorc /////
    md.uorc = uorc_create();
    if (md.uorc == NULL) {
        printf("ERR: Could not create orc\nExiting...\n");
        exit(-1);
    }
    uorc_motor_t motor_right = {.uorc = md.uorc, .port = 1, .invert = 0 };
    uorc_motor_t motor_left  = {.uorc = md.uorc, .port = 0, .invert = 1 };
    uorc_encoder_t encoder_right = {.uorc = md.uorc, .port = 1, .invert = 0};
    uorc_encoder_t encoder_left  = {.uorc = md.uorc, .port = 0, .invert = 1};
    md.motor_right = &motor_right;
    md.motor_left  = &motor_left;
    md.encoder_right = &encoder_right;
    md.encoder_left  = &encoder_left;

    ///// Signal handling /////
	on_exit(cleanup, &md);
	signal(SIGINT, sighandler);
	signal(SIGTERM, sighandler);
	signal(SIGSEGV, sighandler);

    md.msgauto.msg = NULL;
    md.msgencoder.msg = NULL;
    md.msgmanual.msg = NULL;

    ///// Start Threads /////
    pthread_mutex_init(&md.msgauto.mutex, NULL);
    pthread_mutex_init(&md.msgencoder.mutex, NULL);
    pthread_mutex_init(&md.msgmanual.mutex, NULL);
    if (!(md.pids = (pthread_t *) calloc(md.num_threads, sizeof(pthread_t)))) {
        perror("pids");
        exit(1);
    }
    start_threads(&md);

    ///// cleanup /////
    getopt_destroy(gopt);

    ///// handle lcm messages /////
    while (1) {
        if (lcm_handle(md.lcm))
            exit(1);
    }
}
